#include "faiss_hnsw_index.h"
#include <faiss/index_io.h>
#include <filesystem>

FaissHNSWIndex::FaissHNSWIndex(const ConfigParser& config) {
    k_f = config.get_int("k_f", 24);
    efs = config.get_int("efs", 200);
    efc = config.get_int("efc", 100);
#ifdef QUANT
    qbs = config.get_int("qbs", 32);
#endif
    std::string metric_str = config.get_string("metric", "L2");
    if (metric_str == "IP") {
        metric = faiss::METRIC_INNER_PRODUCT;
    } else {
        metric = faiss::METRIC_L2;
    }
    index_path = config.get_string("index_path", "index.faiss");
    save_or_load = config.get_string("save_or_load", "neither");
}

FaissHNSWIndex::~FaissHNSWIndex() {
    index.reset();
}

void FaissHNSWIndex::build(const Datasets* data) {
    if (save_or_load == "load") {
        faiss::Index* loaded_index = faiss::read_index(index_path.c_str());
        auto hnsw_ptr = dynamic_cast<faiss::IndexHNSWFlat*>(loaded_index);
        if (!hnsw_ptr) {
            delete loaded_index;
            throw std::runtime_error("Loaded index is not HNSWFlat");
        }
        index.reset(hnsw_ptr);
    } else {
        index = std::make_unique<faiss::IndexHNSWFlat>(data->d, k_f, metric);
        index->hnsw.efConstruction = efc;
#ifdef QUANT
        index->quant_bits = qbs;
#endif
        index->add(data->nb, data->xb);
        if (save_or_load == "save") {
			std::filesystem::path dir_path = std::filesystem::path(index_path).parent_path();
			if (!dir_path.empty() && !std::filesystem::exists(dir_path)) {
                std::filesystem::create_directories(dir_path);
            }
            faiss::write_index(index.get(), index_path.c_str());
        }
    }
}

void FaissHNSWIndex::prepare_search() {
    hnsw_params.efSearch = efs;
}

void FaissHNSWIndex::search_single(const float* xq, int k, float* distances, int64_t* labels) {
    int n = 1;
    index->search(n, xq, k, distances, labels, &hnsw_params);
}

void FaissHNSWIndex::search_batch(const int nq, const float* xq, int k, float* distances, int64_t* labels) {
    index->search(nq, xq, k, distances, labels, &hnsw_params);
}